import asyncio
import websockets
import sqlite3
import json

# 创建数据库连接和游标
try:
    conn = sqlite3.connect('DataBase/chat-history.db')
    print("Opened database successfully")
except sqlite3.Error as e:
    print(e)
    exit()
cursor = conn.cursor()

# 创建消息记录表和用户表
cursor.execute('''CREATE TABLE IF NOT EXISTS messages
                (id INTEGER PRIMARY KEY AUTOINCREMENT,
                sender TEXT,
                message TEXT,
                room TEXT,
                timestamp DATETIME DEFAULT CURRENT_TIMESTAMP)''')
cursor.execute('''CREATE TABLE IF NOT EXISTS users
                (username TEXT PRIMARY KEY,
                password TEXT)''')
conn.commit()

# 定义房间列表和用户状态
rooms = {}
user_sessions = {}

# 存储历史消息记录
async def store_message(sender, message, room):
    cursor.execute("INSERT INTO messages (sender, message, room) VALUES (?, ?, ?)", (sender, message, room))
    conn.commit()

async def handler(websocket, path):
    global rooms, user_sessions

    try:
        while True:
            message = await websocket.recv()
            data = json.loads(message)
            action = data.get("action")
            
            if action == "register":
                username = data.get("username")
                password = data.get("password")
                if username and password:
                    try:
                        cursor.execute("INSERT INTO users (username, password) VALUES (?, ?)", (username, password))
                        conn.commit()
                        await websocket.send(json.dumps({"type": "register_success", "message": "注册成功，请登录"}))
                    except sqlite3.IntegrityError as e:
                        if 'UNIQUE constraint failed' in str(e):
                            await websocket.send(json.dumps({"type": "error", "message": "用户名已存在"}))
                        else:
                            await websocket.send(json.dumps({"type": "error", "message": "注册失败，请重试"}))
                else:
                    await websocket.send(json.dumps({"type": "error", "message": "用户名或密码不能为空"}))
            elif action == "login":
                username = data.get("username")
                password = data.get("password")
                cursor.execute("SELECT password FROM users WHERE username = ?", (username,))
                record = cursor.fetchone()
                if record and record[0] == password:
                    # 保存用户会话
                    user_sessions[websocket] = {"username": username, "room": "主聊天室"}
                    room = user_sessions[websocket]["room"]
                    
                    # 将新连接的客户端添加到对应房间的集合中
                    if room not in rooms:
                        rooms[room] = set()
                    rooms[room].add(websocket)

                    # 发送欢迎信息并存储
                    #await store_message('系统', '有新人加入，欢迎', room)
                    await websocket.send(json.dumps({"type": "login_success", "username": username}))

                    # 发送历史消息记录
                    cursor.execute("SELECT sender, message FROM messages WHERE room = ? ORDER BY timestamp", (room,))
                    history = cursor.fetchall()
                    for row in history:
                        await websocket.send(json.dumps({"type": "history", "sender": row[0], "message": row[1]}))
                else:
                    await websocket.send(json.dumps({"type": "error", "message": "用户名或密码错误"}))
                    await websocket.close()
                    return
            elif action == "change_password":
                username = user_sessions.get(websocket, {}).get("username")
                old_password = data.get("old_password")
                new_password = data.get("new_password")
                if username and old_password and new_password:
                    cursor.execute("SELECT password FROM users WHERE username = ?", (username,))
                    record = cursor.fetchone()
                    if record and record[0] == old_password:
                        cursor.execute("UPDATE users SET password = ? WHERE username = ?", (new_password, username))
                        conn.commit()
                        await websocket.send(json.dumps({"type": "change_password_success", "message": "密码修改成功"}))
                    else:
                        await websocket.send(json.dumps({"type": "change_password_error", "message": "旧密码不正确"}))
                else:
                    await websocket.send(json.dumps({"type": "change_password_error", "message": "密码修改失败，请重试"}))
            else:
                if websocket in user_sessions:
                    username = user_sessions[websocket]["username"]
                    if action == "message":
                        text = data.get("message")
                        room = user_sessions[websocket]["room"]
                        if text and username:
                            # 将消息广播给同一房间的所有客户端
                            for client in rooms.get(room, set()):
                                await client.send(json.dumps({"type": "message", "sender": username, "message": text}))
                            # 存储消息记录
                            await store_message(username, text, room)
                    elif action == "change_room":
                        new_room = data.get("room")
                        if new_room:
                            # 更新用户的房间信息
                            old_room = user_sessions[websocket]["room"]
                            rooms[old_room].remove(websocket)
                            if not rooms[old_room]:
                                del rooms[old_room]

                            user_sessions[websocket]["room"] = new_room
                            if new_room not in rooms:
                                rooms[new_room] = set()
                            rooms[new_room].add(websocket)

                            # 发送房间切换信息
                            await websocket.send(json.dumps({"type": "system", "message": f"您已加入 {new_room}"}))
                            cursor.execute("SELECT sender, message FROM messages WHERE room = ? ORDER BY timestamp", (new_room,))
                            history = cursor.fetchall()
                            for row in history:
                                await websocket.send(json.dumps({"type": "history", "sender": row[0], "message": row[1]}))
    except Exception as e:
        print(f"处理过程中出现错误: {e}")
    finally:
        # 客户端断开连接时清理
        if websocket in user_sessions:
            room = user_sessions[websocket]["room"]
            rooms[room].remove(websocket)
            if not rooms[room]:
                del rooms[room]
            del user_sessions[websocket]

start_server = websockets.serve(handler, "0.0.0.0", 12345)
if start_server:
    print("WebSocket server started.")

asyncio.get_event_loop().run_until_complete(start_server)
asyncio.get_event_loop().run_forever()
